package com.soup.lab.basic.keyable;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Iterator;
import java.util.List;
import java.util.concurrent.*;

/**
 * 如果发现带有同样 key 的任务正在执行，那么直接使用那个任务的结果。
 * 避免任务重复执行。
 *
 * 要实现该功能的任务必须实现 Keyable 接口。
 *
 * @see Keyable
 */
public class SameKeyFreeExecutorService implements ExecutorService {

    private static final Logger logger = LoggerFactory.getLogger(SameKeyFreeExecutorService.class);

    private final ExecutorService executorService;
    private final ConcurrentMap<String, Future<?>> resultCache;
    private final ScheduledExecutorService scheduledThreadPool;

    private SameKeyFreeExecutorService(ExecutorService executorService) {
        this.executorService = executorService;
        this.resultCache = new ConcurrentHashMap<>();
        this.scheduledThreadPool = Executors.newSingleThreadScheduledExecutor();
    }

    private SameKeyFreeExecutorService(ExecutorService executorService, ConcurrentMap<String, Future<?>> resultCache) {
        this.executorService = executorService;
        this.resultCache = resultCache;
        this.scheduledThreadPool = Executors.newSingleThreadScheduledExecutor();
    }

    public static SameKeyFreeExecutorService wrap(ExecutorService executorService) {
        SameKeyFreeExecutorService result = new SameKeyFreeExecutorService(executorService);
        result.startCleanUpTask();
        return result;
    }

    public static SameKeyFreeExecutorService wrap(ExecutorService executorService, ConcurrentMap<String, Future<?>> resultCache) {
        SameKeyFreeExecutorService result = new SameKeyFreeExecutorService(executorService, resultCache);
        result.startCleanUpTask();
        return result;
    }

    private void startCleanUpTask() {
        scheduledThreadPool.scheduleAtFixedRate(new CleanUpTask(), 1, 2, TimeUnit.SECONDS);
    }

    @Override
    public <T> Future<T> submit(Callable<T> task) {
        if(task instanceof Keyable) {
            return (Future<T>) resultCache.computeIfAbsent(((Keyable) task).getKey(), k -> executorService.submit(task));
        }
        return executorService.submit(task);
    }

    @Override
    public <T> Future<T> submit(Runnable task, T result) {
        if(task instanceof Keyable) {
            return (Future<T>) resultCache.computeIfAbsent(((Keyable) task).getKey(), k -> executorService.submit(task, result));
        }
        return executorService.submit(task, result);
    }

    @Override
    public Future<?> submit(Runnable task) {
        if(task instanceof Keyable) {
            return resultCache.computeIfAbsent(((Keyable) task).getKey(), k -> executorService.submit(task));
        }
        return executorService.submit(task);
    }

    @Override
    public <T> List<Future<T>> invokeAll(Collection<? extends Callable<T>> tasks) throws InterruptedException {
        // 避免因移除元素导致迭代顺序变化
        List<? extends Callable<T>> copy = new ArrayList<>(tasks);
        List<Future<T>> result = new ArrayList<>(tasks.size());
        // 挑选并移除有缓存的任务，使用缓存的结果
        Iterator<? extends Callable<T>> iterator = copy.iterator();
        for (int i = 0; iterator.hasNext(); i++) {
            Callable<T> task = iterator.next();
            if(task instanceof Keyable) {
                String key = ((Keyable) task).getKey();
                Future<?> future = resultCache.get(key);
                if(future != null) {
                    result.set(i, (Future<T>) future);
                    iterator.remove();
                }
            }
        }
        // 如果移除过任务，就对结果进行拼接
        if(copy.size() < tasks.size()) {
            List<Future<T>> futures = executorService.invokeAll(copy);
            for (int i = 0, j = 0; i < result.size(); i++) {
                if(result.get(i) == null) {
                    result.set(i, futures.get(j++));
                }
            }
            return result;
        }
        return executorService.invokeAll(copy);
    }

    @Override
    public <T> List<Future<T>> invokeAll(Collection<? extends Callable<T>> tasks, long timeout, TimeUnit unit) throws InterruptedException {
        return executorService.invokeAll(tasks, timeout, unit);
    }

    @Override
    public <T> T invokeAny(Collection<? extends Callable<T>> tasks) throws InterruptedException, ExecutionException {
        return executorService.invokeAny(tasks);
    }

    @Override
    public <T> T invokeAny(Collection<? extends Callable<T>> tasks, long timeout, TimeUnit unit) throws InterruptedException, ExecutionException, TimeoutException {
        return executorService.invokeAny(tasks, timeout, unit);
    }

    @Override
    public void execute(Runnable command) {
        executorService.execute(command);
    }

    @Override
    public void shutdown() {
        executorService.shutdown();
    }

    @Override
    public List<Runnable> shutdownNow() {
        return executorService.shutdownNow();
    }

    @Override
    public boolean isShutdown() {
        return executorService.isShutdown();
    }

    @Override
    public boolean isTerminated() {
        return executorService.isTerminated();
    }

    @Override
    public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException {
        return executorService.awaitTermination(timeout, unit);
    }



    class CleanUpTask implements Runnable {
        @Override
        public void run() {
            logger.debug("执行清理任务");
            Iterator<Future<?>> iterator = resultCache.values().iterator();
            while (iterator.hasNext()) {
                Future<?> next = iterator.next();
                if(next.isDone()) {
                    iterator.remove();
                    logger.debug("移除future:{}", next);
                }
            }
        }
    }

}
